##__docformat__ = "restructuredtext en"
# ******NOTICE***************
# two_game.py module by Shoibal Chakravarty
#
# You may copy and use this module as you see fit with no
# guarantee implied provided you keep this notice in all copies.
# *****END NOTICE************
"""
This module  contains the functions that simulate a scenario where the three 
regions have a two period planning horizon (2nd period myopic). The three 
regions cooperate to maximize the sum of their two period utility functions.
The green energy deployment in the last period is myopic.
"""

__author__ = "Shoibal Chakravarty (shoibalc@princeton.edu)"
__version__ = "0.1"
__date__ = "2011/August"
__license__ = "BSD"
__all__ = ['threeoptimized','threeoptimizedupdate','threeoptimizedregion']

import numpy as np
from myopic import  regionupdate, worldupdate, cumulativecost
from flat2 import flatten
from utility import U3, U3optimized
from two_game import finalperiodupdate
import scipy.optimize

def threeoptimized(r1,r2,r3,w,com,precision=0.5,init_grid_size=40):
    """
    Function threeoptimized simulates a scenario where all three regions have a two
    period planning horizon. The three regions cooperate to maximize the sum of their two period utility functions. 
    The solution involves solving a non-linear optimization problem to determine the best current 
    period deployments of these three regions. The green energy deployment in the last period is myopic.\n\n
    threeoptimized(r1,r2,r3,...):\n
    r1, r2 and r3 solve the three region two period optimial green energy deployment problem.\n
    """
    r1.descriptor = r2.descriptor = r3.descriptor = 'Optimal solution for three regions with one period look ahead.'
    print "\n\nIn this scenario %s, %s  and %s  cooperate to maximize the sum of their utility functions.\n"% (r1.name, r2.name, r3.name)
    for t in range(w.T-1):
        threeoptimizedupdate(r1,r2,r3,w,com,t,precision,init_grid_size)
        worldupdate(r1,r2,r3,w,com,t)
        print "\nPeriod", t+1,"\n", r1.name, "g", r1.g[t], "e", r1.e[t],"\n", r2.name, "g", r2.g[t], "e", r2.e[t],\
        "\n",r3.name, "g", r3.g[t], "e", r3.e[t], "\ncg", w.cg[t+1],"\n"
# the last period is still myopic
    t = w.T-1
    finalperiodupdate(r1,t)
    finalperiodupdate(r2,t)
    finalperiodupdate(r3,t)
    worldupdate(r1,r2,r3,w,com,t)
    print "\nPeriod", t+1,"\n", r1.name, "g", r1.g[t], "e", r1.e[t],"\n", r2.name, "g", r2.g[t], "e", r2.e[t],\
        "\n",r3.name, "g", r3.g[t], "e", r3.e[t], "\ncg", w.cg[t+1],"\n"
    cumulativecost(r1,r2,r3,w)


def threeoptimizedupdate(ra,rb,rc,w,com,t,precision,init_grid_size):
    """
    Function threeoptimizedupdate solves the three region two period optimization problem 
    for the current period (except for the last period in the simulation). 
    The function successively calls function threeoptimizedregion which return 
    the current best guess and a rectangle bounding the solution. 
    This procedures continues till the desired precision is reached.
    """
    maxa = min(ra.M[t]*ra.E[t+1],(ra.E[t+1] - ra.G[t]))
    mina = 1e-6
    maxb = min(rb.M[t]*rb.E[t+1],(rb.E[t+1] - rb.G[t]))
    minb = 1e-6
    maxc = min(rc.M[t]*rc.E[t+1],(rc.E[t+1] - rc.G[t]))
    minc = 1e-6
    D0 = 21
    D = D0
    EPS = EPSa = EPSb = EPSc = init_grid_size*D
    print 'Search region',ra.name,'min = 0 max =', maxa, rb.name, 'min = 0, max =', maxb, rc.name, 'min = 0, max =',maxc
    flag = True
    num = False
    xbest, ybest, zbest = np.nan, np.nan, np.nan
    while EPS > precision and flag:
# initial run with  init_grid_size. Subsequent runs reduce it by a factor of D
        if num:
            D = np.random.randint(D0+1,D0+5)
        EPSa = EPSa/D
        EPSb = EPSb/D
        EPSc = EPSc/D
        num = True
        print "GRID SIZE", D
        mina, maxa, minb, maxb, minc, maxc, xbest1, ybest1, zbest1, lenxyz = \
                threeoptimizedregion(ra,rb,rc,w,com,t,mina,maxa,minb,maxb,minc,maxc,EPSa,EPSb,EPSc)
        print "Search region update:", ra.name, "min", mina, "max", maxa, "best", xbest1, rb.name, "min", minb, "max", maxb, "best", ybest1, \
                rc.name, "min", minc, "max", maxc, "best", zbest1
        err_best = (xbest - xbest1)**2 + (ybest - ybest1)**2 + (zbest - zbest1)**2
#        if np.sqrt(err_best) <= 0.25*precision:
# this hack prevents infinite looping if we hit a multiple equilibria. we select a particular solution and carry on.
# it also helps in avoiding an extra round of brute force maximization in some cases.
#            if lenxyz == 1:
#            print "1. ATTAINED REQUIRED PRECISION"
#            flag = False
#            else:
#                print "2. MULTIPLE SOLUTIONS DETECTED, PICKING AVERAGE OF DEGENERATE SOLUTIONS"
#                flag = False
#            break
#        else:
        xbest, ybest, zbest = xbest1, ybest1, zbest1
        EPSa = maxa - mina
        EPSb = maxb - minb
        EPSc = maxc - minc
        EPS = max(EPSa, EPSb, EPSc)
    if flag:
        print "3. ATTAINED REQUIRED PRECISION"
        ga = 0.5*(mina + maxa)
        gb = 0.5*(minb + maxb)
        gc = 0.5*(minc + maxc)
    else:
        ga = xbest
        gb = ybest
        gc = zbest
    
# TODO fix function params

    parama = \
    (gb+gc, com.cg0, com.greenpar, ra.B[t], ra.B[t+1], ra.cn[t], ra.co[t], ra.cn[t+1], ra.co[t+1], ra.G[t],\
    rb.G[t]+rc.G[t], w.G[0], ra.e[t], ra.E[t+1], ra.e[t+1], ra.E[t+2], ra.M[t], ra.M[t+1], (1 - com.discount)**com.years, True)
    u, ra.g[t], ra.u1[t], ra.c1[t], ra.g2[t], ra.u2[t], ra.c2[t], ra.m[t] =  U3(ga,*parama)
    regionupdate(ra,t)

    paramb = \
    (gc+ga, com.cg0, com.greenpar, rb.B[t], rb.B[t+1], rb.cn[t], rb.co[t], rb.cn[t+1], rb.co[t+1], rb.G[t],\
    ra.G[t]+rc.G[t], w.G[0], rb.e[t], rb.E[t+1], rb.e[t+1], rb.E[t+2], rb.M[t], rb.M[t+1], (1 - com.discount)**com.years, True)
    u, rb.g[t], rb.u1[t], rb.c1[t], rb.g2[t], rb.u2[t], rb.c2[t], rb.m[t] =  U3(gb,*paramb)
    regionupdate(rb,t)

    paramc = \
    (ga+gb, com.cg0, com.greenpar, rc.B[t], rc.B[t+1], rc.cn[t], rc.co[t], rc.cn[t+1], rc.co[t+1], rc.G[t],\
    ra.G[t]+rb.G[t], w.G[0], rc.e[t], rc.E[t+1], rc.e[t+1], rc.E[t+2], rc.M[t], rc.M[t+1], (1 - com.discount)**com.years, True)
    u, rc.g[t], rc.u1[t], rc.c1[t], rc.g2[t], rc.u2[t], rc.c2[t], rc.m[t] =  U3(gc,*paramc)
    regionupdate(rc,t)


def threeoptimizedregion(ra,rb,rc,w,com,t,mina,maxa,minb,maxb,minc,maxc,EPSa,EPSb,EPSc):
    """
    Function threeoptimizedregion takes among its parameters the coordinates 
    of a cuboidal search region and a grid-size for discretization. It uses a brute 
    force method (scipy.optimize.brute) to numerically obtain the best solution to the 
    optimization problem - maximum of the sum of the three region's two period utility 
    functions. 
    """
    Wa = 3
    lena = np.ceil((maxa - mina)/EPSa) + 1
    lenb = np.ceil((maxb - minb)/EPSb) + 1
    lenc = np.ceil((maxc - minc)/EPSc) + 1
    EPSa = (maxa - mina)/(lena-1)
    EPSb = (maxb - minb)/(lenb-1)
    EPSc = (maxc - minc)/(lenc-1)
    eps = 1e-5
    slice_abc = np.s_[mina: maxa + eps: EPSa, minb: maxb + eps: EPSb, minc: maxc + eps: EPSc]
    mgrid_a = np.mgrid[slice_abc[0]]
    mgrid_b = np.mgrid[slice_abc[1]]
    mgrid_c = np.mgrid[slice_abc[2]]

    parama = \
    [0.0, com.cg0, com.greenpar, ra.B[t], ra.B[t+1], ra.cn[t], ra.co[t], ra.cn[t+1], ra.co[t+1], ra.G[t],\
    rb.G[t]+rc.G[t], w.G[0], ra.e[t], ra.E[t+1], ra.e[t+1], ra.E[t+2], ra.M[t], ra.M[t+1], (1 - com.discount)**com.years, False]

    paramb = \
    [0.0, com.cg0, com.greenpar, rb.B[t], rb.B[t+1], rb.cn[t], rb.co[t], rb.cn[t+1], rb.co[t+1], rb.G[t],\
    ra.G[t]+rc.G[t], w.G[0], rb.e[t], rb.E[t+1], rb.e[t+1], rb.E[t+2], rb.M[t], rb.M[t+1], (1 - com.discount)**com.years, False]

    paramc = \
    [0.0, com.cg0, com.greenpar, rc.B[t], rc.B[t+1], rc.cn[t], rc.co[t], rc.cn[t+1], rc.co[t+1], rc.G[t],\
    ra.G[t]+rb.G[t], w.G[0], rc.e[t], rc.E[t+1], rc.e[t+1], rc.E[t+2], rc.M[t], rc.M[t+1], (1 - com.discount)**com.years, False]

    paramabc = (parama,paramb,paramc)
    xbest, ybest, zbest = scipy.optimize.brute(U3optimized,slice_abc,args=paramabc,full_output=False,finish=None)
    
    xa  = [int(round((xbest - mina)/EPSa))]
    xb  = [int(round((ybest - minb)/EPSb))]
    xc  = [int(round((zbest - minc)/EPSc))]
    xa1 = [xx - Wa > 0 and mgrid_a[xx - Wa].tolist() or mgrid_a[0].tolist() for xx in xa]
    xb1 = [xx - Wa > 0 and mgrid_b[xx - Wa].tolist() or mgrid_b[0].tolist() for xx in xb]
    xc1 = [xx - Wa > 0 and mgrid_c[xx - Wa].tolist() or mgrid_c[0].tolist() for xx in xc]
    xa2 = [xx + Wa < (lena - 1) and mgrid_a[xx + Wa].tolist() or mgrid_a[-1].tolist() for xx in xa]
    xb2 = [xx + Wa < (lenb - 1) and mgrid_b[xx + Wa].tolist() or mgrid_b[-1].tolist() for xx in xb]
    xc2 = [xx + Wa < (lenc - 1) and mgrid_c[xx + Wa].tolist() or mgrid_c[-1].tolist() for xx in xc]
#    print xa1,'\n',xb1,'\n',xc1,'\n',xa2,'\n',xb2,'\n',xc2,'\n'
    mina, maxa, minb, maxb, minc, maxc = min(flatten(xa1)),max(flatten(xa2)),min(flatten(xb1)),max(flatten(xb2)),min(flatten(xc1)),max(flatten(xc2))
#    print mina, maxa, minb, maxb
#    print "min a", mina, "max a", maxa, "min b", minb, "max b", maxb,"min c", minc, "max c", maxc, "best a", xbest, "best b", ybest, "best c", zbest
    return mina, maxa, minb, maxb, minc, maxc, xbest, ybest, zbest, 1


